# -*- coding: utf-8 -*-
"""
General description:
---------------------
This script shows how use the custom component `solph.custom.Link` to build
a simple transshipment model.


Installation requirements:
---------------------------
This example requires the latest version of oemof. Install by:

    pip install oemof
    pip install matplotlib

12.12.2017 - simon.hilpert@uni-flensburg.de
"""
import pandas as pd
import networkx as nx
from matplotlib import pyplot as plt

# solph imports
from oemof.solph import (EnergySystem, Model, Bus, Flow, Source, Sink,
                         custom, Investment)
from oemof.outputlib import processing, views
from oemof.graph import create_nx_graph


def draw_graph(grph, edge_labels=True, node_color='#AFAFAF',
               edge_color='#CFCFCF', plot=True, node_size=2000,
               with_labels=True, arrows=True, layout='neato'):
    """
    Draw a graph. This function will be removed in future versions.

    Parameters
    ----------
    grph : networkxGraph
        A graph to draw.
    edge_labels : boolean
        Use nominal values of flow as edge label
    node_color : dict or string
        Hex color code oder matplotlib color for each node. If string, all
        colors are the same.

    edge_color : string
        Hex color code oder matplotlib color for edge color.

    plot : boolean
        Show matplotlib plot.

    node_size : integer
        Size of nodes.

    with_labels : boolean
        Draw node labels.

    arrows : boolean
        Draw arrows on directed edges. Works only if an optimization_model has
        been passed.
    layout : string
        networkx graph layout, one of: neato, dot, twopi, circo, fdp, sfdp.
    """
    if type(node_color) is dict:
        node_color = [node_color.get(g, '#AFAFAF') for g in grph.nodes()]

    # set drawing options
    options = {
     'prog': 'dot',
     'with_labels': with_labels,
     'node_color': node_color,
     'edge_color': edge_color,
     'node_size': node_size,
     'arrows': arrows
    }

    # draw graph
    pos = nx.drawing.nx_agraph.graphviz_layout(grph, prog=layout)

    nx.draw(grph, pos=pos, **options)

    # add edge labels for all edges
    if edge_labels is True and plt:
        labels = nx.get_edge_attributes(grph, 'weight')
        nx.draw_networkx_edge_labels(grph, pos=pos, edge_labels=labels)

    # show output
    if plot is True:
        plt.show()


datetimeindex = pd.date_range('1/1/2017', periods=2, freq='H')

es = EnergySystem(timeindex=datetimeindex)

b_0 = Bus(label='b_0')

b_1 = Bus(label='b_1')

es.add(b_0, b_1)

es.add(custom.Link(label="line_0",
                   inputs={
                       b_0: Flow(), b_1: Flow()},
                   outputs={
                       b_1: Flow(investment=Investment()),
                       b_0: Flow(investment=Investment())},
                   conversion_factors={
                       (b_0, b_1): 0.95, (b_1, b_0): 0.9}))


es.add(Source(label="gen_0", outputs={
                                b_0: Flow(nominal_value=100,
                                          variable_costs=50)}))

es.add(Source(label="gen_1", outputs={
                                b_1: Flow(nominal_value=100,
                                          variable_costs=50)}))

es.add(Sink(label="load_0", inputs={
                                b_0: Flow(nominal_value=150,
                                          actual_value=[0, 1],
                                          fixed=True)}))

es.add(Sink(label="load_1", inputs={
                                b_1: Flow(nominal_value=150,
                                          actual_value=[1, 0],
                                          fixed=True)}))

m = Model(energysystem=es)

# m.write('transshipment.lp', io_options={'symbolic_solver_labels': True})

m.solve(solver='cbc',
        solve_kwargs={'tee': True, 'keepfiles': False})

m.results()

graph = create_nx_graph(es, m)

draw_graph(graph, plot=True, layout='neato', node_size=3000,
           node_color={
                  'b_0': '#cd3333',
                  'b_1': '#7EC0EE',
                  'b_2': '#eeac7e'})

results = processing.results(m)

print(views.node(results, 'gen_0'))
print(views.node(results, 'gen_1'))

views.node(results, 'line_0')['sequences'].plot(kind='bar')

# look at constraints of Links in the pyomo model LinkBlock
m.LinkBlock.pprint()
